import numpy as np  
from PIL import Image  
  
def calculate_psnr(img1, img2):  
    """  
    计算两幅图像的峰值信噪比（PSNR）。  
      
    :param img1: 原始图像（NumPy数组）  
    :param img2: 需要比较的图像（NumPy数组）  
    :return: 峰值信噪比（PSNR）  
    """  
    # 确保两个图像具有相同的尺寸  
    assert img1.shape == img2.shape, "Input images must have the same dimensions."  
      
    # 计算均方误差（MSE）  
    mse = np.mean((img1 - img2) ** 2)  
      
    # 如果mse为0，则PSNR为无穷大（图像完全相同）  
    if mse == 0:  
        return float('inf')  
      
    # 峰值信噪比（PSNR）  
    # 假设图像的最大像素值为255（对于8位图像）  
    MAX_I = 255.0  
    PSNR = 20 * np.log10(MAX_I / np.sqrt(mse))  
      
    return PSNR  
  
# 加载图像  
def load_image(image_path):  
    img = Image.open(image_path).convert('RGB')  # 确保图像是RGB格式  
    img = np.array(img, dtype=np.float64)  # 将图像转换为NumPy数组，数据类型为float64  
    img = img / 255.0  # 归一化到[0, 1]范围  
    return img  
  
# 示例图像路径  
image_path1 = 'C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\5.png'  
image_path2 = 'C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\ans.png'  
  
# 加载图像  
img1 = load_image(image_path1)  
img2 = load_image(image_path2)  
  
# 确保两张图像尺寸相同  
if img1.shape != img2.shape:  
    print("Images have different dimensions. Cannot calculate PSNR.")  
else:  
    # 计算PSNR  
    psnr = calculate_psnr(img1, img2)  
    print(f"PSNR: {psnr}")